#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2019 crane <crane@crane-pc>
#
# Distributed under terms of the MIT license.

"""
然后窗口再次右移动1位, 从[1:5) 到 [2:6),  nums[1] out, nums[5] in
先看nums[1] 是0, 所以当前窗口中的1个数不变
再看nums[5] 是0, 所以窗口滑动后, 没有增加什么新的可能的1的组合
"""

import random


class Solution:
    def __init__(self, total, allow_span, nums):
        self.total = total
        self.window = allow_span + 1      # [1, 0, 0, 1], 如果span为3, window为4(window中最左和最右距离不超过span)
        self.nums = nums

    def calc(self):
        # 不用考虑 span >= total情况, 下面代码可以处理这种情况
        f = 0

        first_window = self.nums[0:self.window]      # 选择前span个元素, start_idx 为0的窗口
        window_one_count = first_window.count(1)
        f += window_one_count * (window_one_count-1) // 2       # C(n, 2)

        for window_start_idx in range(1, self.total - self.window + 1):
            window_end_idx = window_start_idx + self.window - 1
            assert window_end_idx < self.total

            if self.nums[window_start_idx-1] == 1:
                window_one_count -= 1

            if self.nums[window_end_idx] == 1:
                # print(f, window_start_idx, window_end_idx)
                f += 1 * window_one_count
                # print(window_one_count)
                window_one_count += 1
                # print(window_one_count, f, window_start_idx, window_end_idx)

        t = self.total * (self.total-1) // 2    # C(n, 2): n选2

        return f, t


def brute(total, span, nums):
    t = total * (total-1) // 2
    f = 0

    for start_idx in range(total):
        if nums[start_idx] != 1:
            continue

        for end_idx in range(start_idx+1, total):
            if end_idx - start_idx > span:
                break

            if nums[end_idx] == 1:
                f += 1

    return f, t


def test():
    l = [1, 0, 1, 1, 0, 0, 0, 1]
    s = Solution(8, 3, l)
    assert s.calc() == brute(8, 3, l)

    l = [1, 0, 1, 1, 0, 1, 0, 1]
    s = Solution(8, 3, l)
    assert s.calc() == brute(8, 3, l)

    l = [1, 0, 1, 1, 0, 1, 1, 1]
    s = Solution(8, 3, l)
    assert s.calc() == brute(8, 3, l)

    l = [1, 1, 1, 1, 0, 1, 1, 1]
    s = Solution(8, 3, l)
    assert s.calc() == brute(8, 3, l)

    # l = 50 * [1] + 50 * [0]
    l = 5000 * [1] + 5000 * [0]
    random.shuffle(l)

    s = Solution(len(l), 5, l)
    assert s.calc() == brute(len(l), 5, l)


    # 测试各种随机情况
    for i in range(105):
        random.shuffle(l)
        # print(l)

        s = Solution(len(l), i, l)
        assert s.calc() == brute(len(l), i, l)

def main():
    print("start main")
    test()

if __name__ == "__main__":
    main()
